import cv2
import math
import os


def cut_video_by_video(root_path: str, video_name: str, shape=[0, 0, 0, 0], need_transpose=0,
                       start=0):
    """

        :param start:
        :param shape:
        :param need_transpose:
            0: no need
            1: original direction left
            2: original direction bottom
            3: original direction right
        :param video_name:
        :param dir_name:
        :param output_path:
        :param root_path: 输入image_names 的目录
        :return:
        """
    """

        :param path:
        :return:
        """
    path_split = root_path.split("\\")
    dir_name = path_split[-1]
    output_dir_name = path_split[-2] + '_' + path_split[-1] + "_cut"
    path_split[-1] = output_dir_name
    output_path = "\\".join(path_split)
    if not os.path.exists(output_path):
        os.mkdir(output_path)
    cap = cv2.VideoCapture(os.path.join(root_path, video_name))
    number = cap.get(7)
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    print("width:{}, height:{}, frame number{}".format(
        cap.get(cv2.CAP_PROP_FRAME_WIDTH),
        cap.get(cv2.CAP_PROP_FRAME_HEIGHT),
        number)
    )
    if need_transpose == 1 or need_transpose == 3:
        temp = width
        width = height
        height = temp
    if number >= 180:
        i = 0
        frames = []
        while True:
            i += 1
            ret, img = cap.read()
            # if len(frames) == 2000:
            #     break
            if not ret:
                break
            if need_transpose == 1:
                img = cv2.transpose(img)
                img = cv2.flip(img, 0)
            elif need_transpose == 2:
                img = cv2.flip(img, 0)
            elif need_transpose == 3:
                img = cv2.transpose(img)

            img = img[shape[0]:int(height - shape[1]), shape[2]:int(width - shape[3])]
            img = cv2.resize(img, (80, 80))
            frames.append(img)
            if i % 1000 == 0:
                print(i)
            # if i % 9000 == 0:
            #     print("9000")
            #
            #     start = start + 100
            #     frames = []
        number = len(frames)
        gen_video(output_path, dir_name, video_name, frames, start, number)
    cap.release()
    print('{} have finished'.format(root_path))
    pass


def gen_video(output_path, dir_name, video_name, frames, start, number):
    fps = 30
    # 图片尺寸
    img_size = (80, 80)
    fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
    # face_video_writer = cv2.VideoWriter(face_path, fourcc, fps, img_size)
    length = math.floor((number / 180) * 2)
    type = str(video_name).split(".")[0]
    if type == '0':
        type_name = "Alert"
    elif type == '5':
        type_name = "Low"
    else:
        type_name = "Drowsy"
    for i in range(5, length-5):
        cut_name = dir_name + "-" + "NO" + "_" + str(i + start) + "-" + type_name + ".avi"
        print(cut_name)
        face_video_writer = cv2.VideoWriter(os.path.join(output_path, cut_name), fourcc, fps, img_size)
        for im_f in frames[i * 90:i * 90 + 180]:
            face_video_writer.write(im_f)
        face_video_writer.release()
